// Calculation for superlattice using explicit Hamiltonian constructed
// in the tight binding sense

#include "tbparameters.hpp"


void calculateBands(string kfilename,
		    int nplaq, 
		    double dplaq, 
		    int nlayers1, 
		    int nlayers2, 
		    string mat1,
		    string mat2,
		    string modelname,
		    cx_vec pol);

void calculate(string kfilename,
	       int nplaq, 
	       double dplaq, 
	       int nlayers1, 
	       int nlayers2, 
	       string mat1,
	       string mat2,
	       string modelname,
	       cx_vec pol);


int main(int argc, char *argv[]){
  
  int    nplaq     = 4;      // number of sides of the plaquette operator
  double dplaq     = 1.0e-5; // side length of the plaquette operator
  int    nlayers1  = 10;      // number of layers of material 1
  int    nlayers2  = 1;      // number of layers of material 2
  string mat1      = "gaas";
  string mat2      = "alas";
  string modelname = "klimeck";
  string kfilename = "kfile.dat";

  double pulsewidth = 100;   // fs
  double centerfreq = 2.1;   // eV
  cx_vec pol;
  pol << 1.0 << cx_double(0.0,1.0);

  int j=0;

  if (argc > j+1){
    string arg1=argv[1];
    if (arg1=="help"){
      cout << "usage: \n      ./sl [kfile nlayers1 nlayers2 mat1 mat2 modelname] \n";
      return 0;
    }
    kfilename=string(argv[j]);
  }
  j++;
  if (argc > j+1){
    nlayers1=atoi(argv[j]);
  }
  j++;
  if (argc > j+1){
    nlayers2=atoi(argv[j]);
  }
  string name;
  j++;
  if (argc > j+1){
    mat1=string(argv[j]);
  }
  j++;
  if (argc > j+1){
    mat2=string(argv[j]);
  }
  j++;
  if (argc > j+1){
    modelname=string(argv[5]);
  }

  cout << "\n WELCOME TO superlattice calculation \n"; 
  cout << "\n\n The following are the parameters:\n"; 
  cout << "Layers: "<<mat1<<" "<<nlayers1<<" layers; "
       <<mat2<<" "<<nlayers2<<" layers \n";
  cout << "Model "<<modelname<<"\n";

  calculateBands(kfilename,nplaq,dplaq,nlayers1,
	    nlayers2,mat1,mat2,modelname,pol);

  return 0;
}
void rearrange(cx_mat &Kappa, cx_mat &D){
  // Solutions rising to the right followed by solutions decaying to the right

  uvec indices = sort_index(real(Kappa),1);
  unsigned int j;
  cx_vec t;
  cx_double temp;
  for (int i=0; i < Kappa.n_elem ; i++){
    j = indices(i);
    temp = Kappa(i);
    Kappa(i) = Kappa(j);
    Kappa(j) = temp;
    t = D.col(i);
    D.col(i) = D.col(j);
    D.col(j) = t;
  }
}
void calculateBands(string kfilename,
	       int nplaq, 
	       double dplaq, 
	       int nlayers1, 
	       int nlayers2, 
	       string mat1,
	       string mat2,
	       string modelname,
	       cx_vec pol){
  //-------------------------------------------------------------------
  // load k file
  //-------------------------------------------------------------------
  mat kdata;
  bool loadkfromfile = true;
  kdata.load(kfilename,raw_ascii);
  // Update the total number of kpoints
  int nk = kdata.n_rows;
  mat kvecs(nk,3);
  imat grid(nk,2);
  for (int i=0;i<nk;i++){
    kvecs(i,0) = kdata(i,0);
    kvecs(i,1) = kdata(i,1);
    kvecs(i,2) = kdata(i,2);
    grid(i,0)  = floor(kdata(i,3));
    grid(i,1)  = floor(kdata(i,4));
  }
  /*-------------------------------------------------------------------
    Calculating:
    1. Hamiltonian
    2. Eigenvalues, eigenvectors      
    3. Velocity matrix elements
    4. Berry curvature
    5. Berry curvature eigenvalues 

    Each of these (except eigenvectors) are stored in their respective
    files. Complex quantities are stored with real and imaginary in 
    different files with real and imag attached to their names as prefixes
   */
  //-------------------------------------------------------------------------
  // Dimensions of the Hamiltonian
  double E;
  cx_mat T,VelocityMatrix[2];
  cout << "Dimensions: "<<T.n_rows<<"\n";
  int    fdim;
  //-------------------------------------------------------------------------
  // Declare arrays depending on the dimension
  mat    eigval,Id;
  //-------------------------------------------------------------------------
  // Index for the top valence band
  int nvalence;
  if (modelname == "vogl")
    nvalence=(nlayers1+nlayers2)*4;
  else
    nvalence=(nlayers1+nlayers2)*8;
  cout << "Top valence band at index "<<nvalence-1<<"\n";
  //-------------------------------------------------------------------------
  // The for loop over k points
  double percent_done, last_display=10;
  cx_mat U,T1,T2,T12,T21;
  vec Energy;
  cx_vec Kappa, K;
  cx_mat  D, B,Q;
  vector<cx_vec> F;
  for (int i=0;i<nk;i++){
    // States at k
    superlatticeT(E,modelname,mat1,nlayers1,mat2,nlayers2,kvecs.row(i),
		  T1,T2,T12,T21);
    if ( i == 0 ){
      fdim = T1.n_rows;
      eigval.zeros(fdim,nk);
      Id.eye(fdim,fdim);
      U.zeros(fdim,fdim);
      Energy.zeros(fdim);
    }

    eig_gen(Kappa,D,T2);
    
    // Rising to right solutions, followed by decaying to right
    rearrange(Kappa,D);
    
    eig_gen(K,B,T1);
    
    // Diagonal representation
    // T21 = inv(D) T21 B
    T21 = T21 * B;
    T21 = solve(D,T21);
    T12 = T12 * D;
    T12 = solve(B,T12);
    T1  = diag(K);

    Q   = T21 * T1 * T12;

    int iplus=0;
    while (real(Kappa(iplus))>=0.0){
      iplus++;
    }
    cx_mat Q1, Q2, eigq, V1;
    Q1 = Q.submat(0,0,iplus-1,iplus-1); 
    Q2 = Q.submat(iplus,0,fdim,iplus-1);

    eig_gen(eigq,V1,Q1);
    vector<int> sols(1,-100);
    for (int ii=0; ii < eigq.n_elem; ii++)
      if ( abs(eigq(ii)) < 1.0e-8)
	sols.push_back(ii);

    if (sols[0] != -100)
      {
	cx_vec fL, fR, f(fdim);
	int i1,i2;
	for (int ii=0; ii < sols.size(); ii++){
	  fL = V1.col(sols(ii));
	  fR = Q2 * fL;
	  for (i1 = 0; i1 < fL.n_elem; i1++)
	    f(i1) = fL(i1);
	  for (i2 = 0; i2 < fR.n_elem; i2++)
	    f(i1+1+i2) = fR(i2);
	  F.push_back(f);
	  Energies.push_back(E);
	}
	
      }
    
  }
  // Save eigenvalues
  //  eigval.save("eig.dat",raw_ascii);
}

void calculate(string kfilename,
	       int nplaq, 
	       double dplaq, 
	       int nlayers1, 
	       int nlayers2, 
	       string mat1,
	       string mat2,
	       string modelname,
	       cx_vec pol){

  //-------------------------------------------------------------------
  // Generate Plaquette vectors
  //-------------------------------------------------------------------
  mat dkvecs;
  rowvec dv(3);
  double plaqarea = nplaq*dplaq*dplaq/2*sin(2*pi/nplaq);

  dkvecs.zeros(nplaq,3);

  vec theta = linspace(0,nplaq-1,nplaq);
  theta = theta*2*pi/(nplaq-1);
  theta = theta - theta(1)/2;
  cout << "Plaquette angles: \n"<< theta/pi*180;

  for (int i=0;i<nplaq;i++){
    dv << cos(theta(i)) << sin(theta(i)) << 0.0;
    dkvecs.row(i) = dplaq*dv;
  }
  //-------------------------------------------------------------------
  // load k file
  //-------------------------------------------------------------------
  mat kdata;
  bool loadkfromfile = true;
  kdata.load(kfilename,raw_ascii);
  // Update the total number of kpoints
  int nk = kdata.n_rows;
  mat kvecs(nk,3);
  imat grid(nk,2);
  for (int i=0;i<nk;i++){
    kvecs(i,0) = kdata(i,0);
    kvecs(i,1) = kdata(i,1);
    kvecs(i,2) = kdata(i,2);
    grid(i,0)  = floor(kdata(i,3));
    grid(i,1)  = floor(kdata(i,4));
  }
  


  /*-------------------------------------------------------------------
    Calculating:
    1. Hamiltonian
    2. Eigenvalues, eigenvectors      
    3. Velocity matrix elements
    4. Berry curvature
    5. Berry curvature eigenvalues 

    Each of these (except eigenvectors) are stored in their respective
    files. Complex quantities are stored with real and imaginary in 
    different files with real and imag attached to their names as prefixes
   */

  //-------------------------------------------------------------------------
  // Dimensions of the Hamiltonian
  double E;
  cx_mat T,VelocityMatrix[2];
  superlatticeH(E,modelname,mat1,nlayers1,
		mat2,nlayers2,
		kvecs.row(0),T,false,0);
  cout << "Dimensions: "<<T.n_rows<<"\n";
  int    fdim=T.n_rows;
  //-------------------------------------------------------------------------
  // Declare arrays depending on the dimension
  mat    eigval(fdim,nk);
  vec    eigvals[nplaq];
  cx_mat eigvecs[nplaq];
  mat Id(fdim,fdim);
  Id.eye();

  cx_mat B[nk];
  cx_mat L[nplaq];

  cx_mat y1;
  mat be(nk,fdim);
  vec b1;
  cx_mat P(fdim,fdim);
  mat bnet(nk,fdim);
  //-------------------------------------------------------------------------
  // Index for the top valence band
  int nvalence;
  if (modelname == "vogl")
    nvalence=(nlayers1+nlayers2)*4;
  else
    nvalence=(nlayers1+nlayers2)*8;
  cout << "Top valence band at index "<<nvalence-1<<"\n";
  //-------------------------------------------------------------------------
  // Indices for the groups of bands: [start, end]
  imat ind(fdim/2,2);
  int  ii=0;
  for (int i=0;i<fdim; i += 2 ){
    ind(ii,0) = i;
    ind(ii,1) = i+1;
    ii++;
  }
  cout << "Index ranges for groups of bands:\n" <<ind<<"\n";
  //-------------------------------------------------------------------------
  // Generate the Mask to act as projectors by elementwise multiplication
  int l,p;
  imat mask(fdim,fdim);
  mask.zeros(fdim,fdim);
  vec *b11;
  b11 = new vec[ind.n_rows];
  for (int i=0;i<ind.n_rows;i++){
    b11[i].zeros(ind(i,1)-ind(i,0)+1);
    for (l=ind(i,0);l<=ind(i,1);l++)
      for (p=ind(i,0);p<=ind(i,1);p++){
	mask(l,p) = 1;
      }
  }
  //-------------------------------------------------------------------------
  // The for loop over k points
  double percent_done, last_display=10;

  // Gauge transform matrix
  cx_mat *W, *Hamiltonian, *HThz, *Rho;
  W = new cx_mat[nk];
  Rho = new cx_mat[nk];
  HThz = new cx_mat[nk];
  Hamiltonian = new cx_mat[nk];
  W[0].eye(fdim,fdim);
  cx_double cmplxi(0.0,1.0);
  cx_mat U(fdim,fdim);
  vec Energy(fdim);
  cout << "cmplxi "<<cmplxi<<"\n";
  for (int i=0;i<nk;i++){
    // States at k
    superlatticeH(E,modelname,mat1,nlayers1,
		  mat2,nlayers2,
		  kvecs.row(i),T,false,0);
    eig_sym(eigvals[0],eigvecs[0],T);
    eigval.col(i)=eigvals[0];
    
    //----------------------------------------------------------------------
    // Velocity matrix at the first point of the plaquette. Representation
    // in eigenvectors is done below with respect to the first point on
    // the plaquette. 
    //
    // Units of velocity matrix = Units of dH/dk  = E x Length
    for (int mu=0;mu<2;mu++){
      // run superlatticeH in d/dk mode
      superlatticeH(E,modelname,mat1,nlayers1,
		    mat2,nlayers2,kvecs.row(i)+dkvecs.row(0),
		    VelocityMatrix[mu],true,mu);
      // multiply by length scale:
      
    }
    //----------------------------------------------------------------------         
    // Berry
    for (int j=0;j<nplaq;j++){
      superlatticeH(E,modelname,mat1,nlayers1,
		    mat2,nlayers2,
		    kvecs.row(i)+dkvecs.row(j),T,false,0);
      eig_sym(eigvals[j],eigvecs[j],T);
      if (j==0)
	Hamiltonian[i] = T;
    }
    Energy = eigvals[0];
    B[i]=berry(eigvecs,mask,nplaq,plaqarea,L);
    //B[i] = berry_sumrule(VelocityMatrix[0],VelocityMatrix[1],mask,eigvals[0]);
    for (int ii=0;ii<ind.n_rows;ii++){
      l = ind(ii,0);
      p = ind(ii,1);
      if ( p < nvalence ){
	eig_sym(b11[ii],y1,B[i].submat(l,l,p,p));
	b11[ii] = sort(b11[ii]);
	be.submat(i,l,i,p)=b11[ii].t();
      }
    }
    //----------------------------------------------------------------------
    // HThz units =  E x Length / Energy = Length
    // HThz x (e Electric field) = Length x Energy/Length = Energy
    HThz[i] = VelocityMatrix[0]*(1-mask);
    for (int j=0;j<fdim;j++)
      for (int ii=0;ii<fdim;ii++)
	if (mask(ii,j) == 0 )
	  HThz[i](ii,j) = -cmplxi*VelocityMatrix[0](ii,j)
	                         /(eigvals[0](ii)-eigvals[0](j));

    //----------------------------------------------------------------------
    // Gauge transformation matrix: 
    // L[0] = 1st point of the plaquette
    // For a square plaquette, L[0] = x link. Keep it this way!

    if (i < nk-1)
      if (grid(i+1,0) == 0)
	W[i+1].eye(fdim,fdim);
      else
	W[i+1] = W[i] * L[0];
    //----------------------------------------------------------------------
    // Apply Gauge transformations
    U = eigvecs[0]*W[i].t();
    for (int mu=0;mu<2;mu++){
      VelocityMatrix[mu]  =  U.t()*VelocityMatrix[mu]*U;
    }
    B[i]                  =  W[i] *B[i] *W[i].t();
    // Note: Because W is allowed to mix spin split bands, Hamiltonian
    //       is not expected to be diagonal
    Hamiltonian[i]        =  U.t()*Hamiltonian[i]*U;
    HThz[i]               =  U.t()*HThz[i]*U;
    //----------------------------------------------------------------------
    // Optical transition from the valence to conduction via Vx + i Vy
    P = VelocityMatrix[0] - cx_double(0.0,1.0)*VelocityMatrix[1];
    P.submat(0,0,nvalence-1,nvalence-1)       *= 0;
    P.submat(nvalence,nvalence,fdim-1,fdim-1) *= 0;

    for (l = 0; l < nvalence ; l++ )
      for (p = nvalence; p < fdim ; p ++){
	P(l,p) =  cmplxi*P(l,p) / (Energy(l)-Energy(p));
	P(p,l) = -cmplxi*P(p,l) / (Energy(p)-Energy(l));
      }
    P = P * P.t();
    

    // Initial Density matrix =  Rho[i](l,l') (|F(w - wg - wl)|^2+|F(w - wg - wl')|^2)/2
    //       
    //       where F(w) is the Fourier transform of Electric field
    //
    Rho[i] = P.submat(0,0,nvalence-1,nvalence-1);

    //----------------------------------------------------------------------
    // Net Berry
    U = P * B[i];
    bnet.row(i) = real(U.diag()).t();

    //----------------------------------------------------------------------
    // Progress indicator
    percent_done = i*100.0/nk;
    if (percent_done >= last_display + 10){
      last_display = percent_done;
      cout << (i*100)/nk <<" % complete \n";
    }
  }
 

  cout << "complete. writing files...\n";

  mat t;
  // Save eigenvalues
  eigval.save("eig.dat",raw_ascii);
  be.save("rbe.dat",raw_ascii);
  bnet.save("bnet.dat",raw_ascii);

  cout << "FINISHED\n";

  return;
}


